package ltcode

import (
	"crypto/rand"
	"fmt"
	"math/big"
)

// LTEncoder implements Luby Transform encoding
type LTEncoder struct {
	dataShards int
	totalShards int
	data [][]byte
	shardSize int
}

// LTDecoder implements Luby Transform decoding
type LTDecoder struct {
	dataShards int
	totalShards int
	received [][]byte
	receivedCount int
	shardSize int
}

// NewLTEncoder creates a new LT encoder
func NewLTEncoder(dataShards, totalShards int) *LTEncoder {
	return &LTEncoder{
		dataShards: dataShards,
		totalShards: totalShards,
		data: make([][]byte, totalShards),
	}
}

// NewLTDecoder creates a new LT decoder
func NewLTDecoder(dataShards, totalShards int) *LTDecoder {
	return &LTDecoder{
		dataShards: dataShards,
		totalShards: totalShards,
		received: make([][]byte, totalShards),
		receivedCount: 0,
	}
}

// Encode encodes data using LT code
func (enc *LTEncoder) Encode(input []byte) ([][]byte, error) {
	if len(input) == 0 {
		return nil, fmt.Errorf("input data is empty")
	}

	// Calculate shard size
	enc.shardSize = (len(input) + enc.dataShards - 1) / enc.dataShards
	
	// Pad input to make it divisible by dataShards
	paddedInput := make([]byte, enc.dataShards * enc.shardSize)
	copy(paddedInput, input)
	
	// Split input into data shards
	for i := 0; i < enc.dataShards; i++ {
		start := i * enc.shardSize
		end := start + enc.shardSize
		if end > len(paddedInput) {
			end = len(paddedInput)
		}
		enc.data[i] = make([]byte, enc.shardSize)
		copy(enc.data[i], paddedInput[start:end])
	}

	// Generate LT code symbols
	ltSymbols := make([][]byte, enc.totalShards)
	
	for i := 0; i < enc.totalShards; i++ {
		// Generate random degree for this symbol
		degree := enc.generateDegree()
		
		// XOR selected data shards
		symbol := make([]byte, enc.shardSize)
		selectedShards := enc.selectShards(degree)
		
		for _, shardIndex := range selectedShards {
			if shardIndex < enc.dataShards {
				for j := 0; j < enc.shardSize; j++ {
					symbol[j] ^= enc.data[shardIndex][j]
				}
			}
		}
		
		// Add minimal metadata: only degree (use 1 byte for degree <= 255)
		metadata := make([]byte, 1 + len(selectedShards))
		metadata[0] = byte(degree)
		for j, shardIndex := range selectedShards {
			metadata[1+j] = byte(shardIndex)
		}
		
		// Combine metadata and symbol
		ltSymbols[i] = append(metadata, symbol...)
	}
	
	return ltSymbols, nil
}

// Decode decodes LT code symbols
func (dec *LTDecoder) Decode(symbols [][]byte) ([]byte, error) {
	if len(symbols) == 0 {
		return nil, fmt.Errorf("no symbols to decode")
	}

	// Extract metadata and symbols
	dec.shardSize = len(symbols[0]) - 4 // Assume all symbols have same size
	decoded := make([][]byte, dec.dataShards)
	decodedCount := 0
	
	// Initialize decoded shards
	for i := 0; i < dec.dataShards; i++ {
		decoded[i] = make([]byte, dec.shardSize)
	}
	
	// Process symbols
	for _, symbol := range symbols {
		if len(symbol) < 1 {
			continue
		}
		
		// Extract degree
		degree := int(symbol[0])
		
		// Extract selected shards
		selectedShards := make([]int, degree)
		for i := 0; i < degree; i++ {
			selectedShards[i] = int(symbol[1+i])
		}
		
		// Extract symbol data
		symbolData := symbol[1+degree:]
		
		// If degree is 1, we can directly decode
		if degree == 1 && selectedShards[0] < dec.dataShards {
			decoded[selectedShards[0]] = make([]byte, dec.shardSize)
			copy(decoded[selectedShards[0]], symbolData)
			decodedCount++
		}
	}
	
	// If we have enough decoded shards, reconstruct original data
	if decodedCount >= dec.dataShards {
		result := make([]byte, 0, dec.dataShards * dec.shardSize)
		for i := 0; i < dec.dataShards; i++ {
			result = append(result, decoded[i]...)
		}
		return result, nil
	}
	
	return nil, fmt.Errorf("insufficient symbols for decoding")
}

// generateDegree generates a random degree for LT code
func (enc *LTEncoder) generateDegree() int {
	// Use ideal soliton distribution
	// For efficiency, limit max degree to 5
	maxDegree := 5
	if enc.dataShards < maxDegree {
		maxDegree = enc.dataShards
	}
	
	// Generate random degree
	degree, _ := rand.Int(rand.Reader, big.NewInt(int64(maxDegree)))
	return int(degree.Int64()) + 1
}

// selectShards randomly selects shards for XOR
func (enc *LTEncoder) selectShards(degree int) []int {
	shards := make([]int, degree)
	
	for i := 0; i < degree; i++ {
		// Generate random shard index
		shardIndex, _ := rand.Int(rand.Reader, big.NewInt(int64(enc.dataShards)))
		shards[i] = int(shardIndex.Int64())
	}
	
	return shards
}

// Verify verifies the integrity of encoded data
func (enc *LTEncoder) Verify(data [][]byte) (bool, error) {
	// Simple verification - check if we have enough symbols
	if len(data) < enc.dataShards {
		return false, fmt.Errorf("insufficient symbols for verification")
	}
	
	// For LT code, we can't easily verify without decoding
	// This is a simplified implementation
	return true, nil
}

// Reconstruct reconstructs missing data shards
func (dec *LTDecoder) Reconstruct(data [][]byte) error {
	// This is a simplified implementation
	// In practice, LT code reconstruction is more complex
	return nil
}
